# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Space updater"""

from typing import TYPE_CHECKING, Union, List, Optional, Tuple
from numpy import ndarray

import torch
from torch import Tensor
from torch.nn import Parameter, ModuleList
from torch.optim import Optimizer

from..system import Molecule
from ..control import Controller
from ..function import get_tensor, get_arguments
from ..function import functions as func


class Updater(Optimizer):
    r"""
    Base class of the MindSPONGE updater, which is a special subclass of the `Optimizer` in PyTorch.
    The `Updater` updates the atomic coordinates of the simulation system. The updating of atomic coordinates
    requires atomic forces and atomic velocities, where the force is passed from outside and the velocity is the
    parameter of the `Updater` itself. And in the case of periodic boundary conditions (PBC), the `Updater`
    could also update the size of the PBC box by the virial of the simulation system.
    The "Updater" controls the values of seven variables during the simulation through a series of `Controller`:
    coordinates, velocity, force, energy, kinetics, virial and pbc_box. If more than one `Controller` is passed in,
    they will work in sequence.

    Args:
        system(Molecule):                                   Simulation system.
        controller(Union[Controller, List[Controller]]):    Controller or list of controllers to control the seven
                                                            variables (coordinate, velocity, force, energy, kinetics,
                                                            virial and pbc_box) of the simulation system. Default: None
        time_step(float):                                   Time step. Defulat: 1e-3
        velocity(Union[Tensor, ndarray, List[float]]):      Array of atomic velocity. The shape of array is `(A, D)`
                                                            or `(B, A, D)`, and the data type is float. Default: None
        weight_decay(float):                                An value for the weight decay. Default: 0.0
        loss_scale(float):                                  A value for the loss scale. Default: 1.0

    Inputs:
        - **energy** (Tensor) - Energy of the system. Tensor of shape `(B, A, D)`. Data type is float.
        - **force** (Tensor) - Force of the system. Tensor of shape `(B, A, D)`. Data type is float.
        - **virial** (Tensor) - Virial of the system. Tensor of shape `(B, A, D)`. Data type is float. Default: None

    Outputs:
        bool, whether successfully finish the current optimization step and move to next step.

    Supported Platforms:
        ``Ascend`` ``GPU``

    Symbols:
        B:  Batchsize, i.e. number of walkers in simulation
        A:  Number of atoms.
        D:  Spatial dimension of the simulation system. Usually is 3.
    """
    def __init__(self,
                 system: Molecule,
                 controller: Union[Controller, List[Controller]] = None,
                 time_step: float = 1e-3,
                 velocity: Union[Tensor, ndarray, List[float]] = None,
                 weight_decay: float = 0.0,
                 loss_scale: float = 1.0,
                 **kwargs
                 ):
        # 获取系统参数
        params = system.space_parameters()
        
        # 调用父类初始化
        defaults = dict(lr=time_step, weight_decay=weight_decay, loss_scale=loss_scale)
        super().__init__(params, defaults)

        self._kwargs = get_arguments(locals(), kwargs)
        self._kwargs.pop('velocity')

        # 基本参数设置
        self.time_step = get_tensor(time_step, dtype=torch.float32)
        self.system = system
        
        # 系统状态
        self.coordinate = self.system.coordinate
        self.pbc_box = self.system.pbc_box
        
        # 当前状态记录
        self.current_coordinate = self.coordinate.detach().clone()
        self.current_pbc_box = None if self.pbc_box is None else self.pbc_box.detach().clone()

        # 原子质量相关
        self.atom_mass = self.system.atom_mass
        self.inv_mass = self.system.inv_mass
        # (B,A,1)
        self._atom_mass = self.atom_mass.unsqueeze(-1)
        self._inv_mass = self.inv_mass.unsqueeze(-1)

        # 系统维度信息
        self.num_walker = system.num_walker
        self.num_atoms = system.num_atoms
        self.dimension = system.dimension

        # 单位转换
        self.units = self.system.units
        self.kinetic_unit_scale = get_tensor(self.units.kinetic_ref, dtype=torch.float32)

        # 初始化速度 - 使用普通Tensor
        if velocity is None:
            self.velocity = torch.zeros_like(self.coordinate)
        else:
            velocity = get_tensor(velocity, dtype=torch.float32)
            if velocity.ndim == 2:
                velocity = velocity.unsqueeze(0)
            if velocity.shape != self.coordinate.shape:
                raise ValueError(f'Velocity shape {velocity.shape} must match coordinate shape {self.coordinate.shape}')
            self.velocity = velocity

        # 约束和控制器设置
        self.num_constraints = 0
        self.num_controller = 0
        self.controller = None
        if controller is not None:
            if isinstance(controller, Controller):
                self.num_controller = 1
                controller = [controller]
            elif isinstance(controller, list):
                self.num_controller = len(controller)
            else:
                raise TypeError(f'controller must be Controller or list but got: {type(controller)}')

            self.controller = ModuleList(controller)
            for i in range(self.num_controller):
                self.num_constraints += self.controller[i].num_constraints

        # 自由度设置
        self.sys_dofs = system.degrees_of_freedom
        self.degrees_of_freedom = 0
        self.set_degrees_of_freedom(self.sys_dofs - self.num_constraints)

        # 初始化受力相关参数 - 使用普通Tensor
        self.potential = get_tensor(torch.zeros((self.coordinate.shape[0], 1), dtype=torch.float32))
        self.force = None
        if self.velocity is not None:
            self.force = torch.zeros_like(self.coordinate)

        # 初始化动力学参数 - 使用普通Tensor
        self.kinetics = None
        self.temperature = None
        if self.velocity is not None:
            kinetics = self.get_kinetics(self.velocity)
            temperature = self.get_temperature(kinetics)
            self.kinetics = kinetics
            self.temperature = temperature

        # 初始化压强相关参数 - 使用普通Tensor
        self.virial = None
        self.pressure = None
        if self.pbc_box is not None:
            self.virial = get_tensor(torch.zeros((self.num_walker, self.dimension), dtype=torch.float32))
            self.pressure = get_tensor(torch.zeros((self.num_walker, self.dimension), dtype=torch.float32))

        # 步数记录 - 使用普通Tensor
        self.steps = get_tensor(0, dtype=torch.int32)

        # 设置控制器时间步长
        if controller is not None:
            for i in range(self.num_controller):
                self.controller[i].set_time_step(self.time_step)

    @property
    def boltzmann(self) -> float:
        """
        Boltzmann constant in current unit.

        Returns:
            float, Boltzmann constant in current unit.
        """
        return self.units.boltzmann

    @property
    def press_unit_scale(self) -> float:
        """
        Reference value of pressure.

        Returns:
            float, reference value of pressure.
        """
        return self.units.pressure_ref

    def set_steps(self, steps: int = 0):
        """
        Set current step of the system.

        Args:
            step(int):  Current step of the system. Default: 0
        """
        steps = get_tensor(steps, dtype=torch.int32)
        self.steps = steps
        return self

    def set_degrees_of_freedom(self, dofs: int):
        """
        Set degrees of freedom (DOFs)

        Args:
            dofs(int):  Degrees of freedom.
        """
        self.degrees_of_freedom = func.get_integer(dofs)
        self.num_constraints = self.sys_dofs - self.degrees_of_freedom
        for i in range(self.num_controller):
            self.controller[i].set_degrees_of_freedom(self.degrees_of_freedom)
        return self

    def update_coordinate(self, coordinate: Tensor, success: bool = True) -> bool:
        """Update the coordinate tensor"""
        self.system.update_coordinate(coordinate)
        return success

    def update_pbc_box(self, pbc_box: Tensor, success: bool = True) -> bool:
        """Update the PBC box tensor"""
        if self.system.pbc_box is not None:
            self.system.update_pbc_box(pbc_box)
        return success

    def update_velocity(self, velocity: Tensor, success: bool = True) -> bool:
        """Update the velocity tensor"""
        self.velocity.data.copy_(velocity)
        return success

    def record_coordinate(self, coordinate: Tensor, success: bool = True) -> bool:
        """Record current coordinate"""
        self.current_coordinate.copy_(coordinate.detach().clone())
        return success

    def record_pbc_box(self, pbc_box: Tensor, success: bool = True) -> bool:
        """Record current PBC box"""
        if self.current_pbc_box is not None:
            self.current_pbc_box.copy_(pbc_box.detach().clone())
        return success

    def record_kinetics(self, kinetics: Tensor, success: bool = True) -> bool:
        """Record current kinetics"""
        if self.kinetics is None:
            self.kinetics = kinetics.clone()
        else:
            self.kinetics.copy_(kinetics)
        return success

    def record_temperature(self, temperature: Tensor, success: bool = True) -> bool:
        """Record current temperature"""
        if self.temperature is None:
            self.temperature = temperature.clone()
        else:
            self.temperature.copy_(temperature)
        return success

    def record_virial(self, virial: Tensor, success: bool = True) -> bool:
        """Record current virial"""
        if self.virial is not None:
            self.virial.copy_(virial)
        return success

    def record_pressure(self, pressure: Tensor, success: bool = True) -> bool:
        """Record current pressure"""
        if self.pressure is not None:
            self.pressure.copy_(pressure)
        return success

    def get_velocity(self) -> Tensor:
        """
        Get velocity.

        Returns:
            Tensor, atom velocities of the system.
        """
        return self.velocity if self.velocity is not None else None

    def get_kinetics(self, velocity: Tensor) -> Tensor:
        """
        Get kinectics.

        Args:
            velocity(Tensor):   Tensor of atom velocities. Data type is float.

        Returns:
            Tensor, the kinectics of the system.
        """
        # (B,A,D)
        kinetics = 0.5 * self._atom_mass * velocity**2
        # (B,D) <- (B,A,D)
        kinetics = kinetics.sum(-2)
        return kinetics * self.kinetic_unit_scale

    def get_temperature(self, kinetics: Tensor = None) -> Tensor:
        """
        Get temperature.

        Args:
            kinetics(Tensor):   Tensor of kinetics. Data type is float. Default: None

        Returns:
            Tensor, the temperature of the system.
        """
        # (B) <- (B,D)
        kinetics = kinetics.sum(-1)
        return 2 * kinetics / self.degrees_of_freedom / self.boltzmann

    def get_pressure(self, kinetics: Tensor, virial: Tensor, pbc_box: Tensor) -> Tensor:
        """
        Get pressure.

        Args:
            kinetics(Tensor):   Tensor of kinetics. Data type is float.
            virial(Tensor):     Tensor of virial. Data type is float.
            pbc_box(Tensor):    Tensor of pbc_box. Data type is float.

        Returns:
            Tensor, the pressure of the system.
        """
        if self.pbc_box is None:
            return None
        # (B,D) = ((B,D) - (B, D)) / (B,1)
        volume = torch.prod(pbc_box, -1, keepdim=True)
        pressure = 2 * (kinetics - virial) / volume
        return pressure * self.press_unit_scale

    def get_dt(self):
        """
        Get the learning rate of current step.

        Returns:
            float, the learning rate of current step.
        """
        return self.defaults['lr']

    def next_step(self) -> bool:
        """
        Finish the current optimization step and move to next step.

        Args:
            success(bool):  Whether to finish the current optimization step and move to next step. Default: True

        Returns:
            bool, whether successfully finish the current optimization step and move to next step.
        """
        self.steps.data.copy_(self.steps.data + 1)

    def decay_and_scale_grad(self, force: Tensor, virial: Tensor = None) -> Tuple[Tensor, Tensor]:
        """
        Do weight decay and gradient scale for force and virial.

        Args:
            force(Tensor):  Tensor of force. Data type is float.
            virial(Tensor): Tensor of virial. Data type is float. Default: None

        Returns:
            - Tensor, Tensor of force after weight decay and gradient scale.
            - Tensor, Tensor of virial after weight decay and gradient scale.
              If pbc_box is None, the output virial is the same as input.
        """
        if self.defaults['weight_decay'] or self.defaults['loss_scale']:
            if self.pbc_box is None:
                gradients = (force,)
            else:
                gradients = (force, virial)

            gradients = self.weight_decay(gradients)
            gradients = self.scale_grad(gradients)

            force = gradients[0]
            if self.pbc_box is not None:
                virial = gradients[1]

        return force, virial
    def weight_decay(self, gradients: Tuple[Tensor, ...]) -> Tuple[Tensor, ...]:
        """Apply weight decay to gradients."""
        if self.defaults['weight_decay'] <= 0:
            return gradients
        
        decayed_grads = []
        for grad in gradients:
            if grad is not None:
                decayed_grads.append(grad + self.defaults['weight_decay'] * grad)
            else:
                decayed_grads.append(None)
        return tuple(decayed_grads)

    def scale_grad(self, gradients: Tuple[Tensor, ...]) -> Tuple[Tensor, ...]:
        """Scale gradients by loss scale."""
        if self.defaults['loss_scale'] == 1.0:
            return gradients
        
        scaled_grads = []
        for grad in gradients:
            if grad is not None:
                scaled_grads.append(grad / self.defaults['loss_scale'])
            else:
                scaled_grads.append(None)
        return tuple(scaled_grads)

    def step(self, closure):
        """执行一步优化更新
        
        Args:
            closure (callable, optional): 用于重新评估模型的闭包函数
            
        Returns:
            bool: 是否成功完成当前步骤
        """
        # 执行closure获取能量
        energy = None
        if closure is not None:
            energy = closure()
            
        # 从梯度获取力和维里
        force = -self.system.coordinate.grad
        virial = None
        if self.pbc_box is not None and self.pbc_box.grad is not None:
            virial = 0.5 * self.pbc_box.grad * self.pbc_box
            
        # 执行更新
        force, virial = self.decay_and_scale_grad(force, virial)

        coordinate = self.coordinate
        velocity = self.velocity
        pbc_box = self.pbc_box

        variables = {
            'coordinate': coordinate,
            'velocity': velocity,
            'energy': energy,
            'force': force,
            'virial': virial,
            'pbc_box': pbc_box,
        }

        # 应用控制器
        if self.controller is not None:
            for controller in self.controller:
                variables = controller(**variables, steps=self.steps)

        # 更新动力学量
        kinetics = self.get_kinetics(variables['velocity'])
        temperature = self.get_temperature(kinetics)
        pressure = self.get_pressure(kinetics, virial, pbc_box) if pbc_box is not None else None

        # 记录状态
        self.record_coordinate(variables['coordinate'])
        self.record_pbc_box(variables['pbc_box'])
        self.record_kinetics(kinetics)
        self.record_virial(virial)
        self.record_temperature(temperature)
        self.record_pressure(pressure)

        # 更新参数
        self.update_coordinate(variables['coordinate'])
        self.update_velocity(variables['velocity'])
        self.update_pbc_box(variables['pbc_box'])

        # 更新步数并返回
        return self.next_step()

    def zero_grad(self, set_to_none: bool = False):
        """Clears the gradients of all optimized parameters."""
        for group in self.param_groups:
            for p in group['params']:
                if set_to_none:
                    p.grad = None
                elif p.grad is not None:
                    p.grad.detach_()
                    p.grad.zero_()
                else:
                    p.grad = torch.zeros_like(p)

    def state_dict(self):
        """Returns the state of the optimizer as a dict."""
        state_dict = super().state_dict()
        state_dict.update({
            'steps': self.steps,
            'potential': self.potential,
            'force': self.force,
            'velocity': self.velocity,
            'kinetics': self.kinetics,
            'temperature': self.temperature,
            'virial': self.virial,
            'pressure': self.pressure,
            'current_coordinate': self.current_coordinate,
            'current_pbc_box': self.current_pbc_box,
        })
        return state_dict

    def load_state_dict(self, state_dict):
        """Loads the optimizer state."""
        super().load_state_dict(state_dict)
        self.steps.data.copy_(state_dict['steps'])
        self.potential.data.copy_(state_dict['potential'])
        self.force.data.copy_(state_dict['force'])
        self.velocity.data.copy_(state_dict['velocity'])
        self.kinetics.data.copy_(state_dict['kinetics'])
        self.temperature.data.copy_(state_dict['temperature'])
        if self.virial is not None:
            self.virial.data.copy_(state_dict['virial'])
        if self.pressure is not None:
            self.pressure.data.copy_(state_dict['pressure'])
        self.current_coordinate.data.copy_(state_dict['current_coordinate'])
        if self.current_pbc_box is not None:
            self.current_pbc_box.data.copy_(state_dict['current_pbc_box'])

